from define_problem_space_parameters import number_of_time_steps
from initialize_updating_coefficients import Chxh,Chxey,Chxez,Chyh,Chyez,Chyex,Chzex,Chzey,Chzh

from initialize_boundary_conditions import *

import time
import numpy as np

print ('Starting the time marching loop')
print(f'Total number of time steps : {number_of_time_steps}') 

start_time = time.process_time()
current_time = 0.0
def printInfo(fields,i,j,k):
    print("%.30f"%(fields[i-1,j-1,k-1]))
def logInfo(time_step_beg,time_step_end,strLabel,fields,i,j,k):
    if(time_step+1>=time_step_beg and time_step+1<=time_step_end):
        print("%s:%.30f"%(strLabel,fields[i-1,j-1,k-1]))
def debug_pause(time_step_beg,time_step_end,strLabel):
    if(time_step+1>=time_step_beg and time_step+1<=time_step_end):
        print("%s"%(strLabel))
        time.sleep(1000)
def update_magnetic_fields():
    from run_fdtd_time_marching_loop import current_time
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez
    current_time  = current_time + dt/2
    # Hx update
    Hx[:,:,:] = (Chxh * Hx[:,:,:] + 
        Chxey * (Ey[0:nxp1, 0:ny, 1:nzp1] - Ey[0:nxp1, 0:ny, 0:nz]) + 
        Chxez * (Ez[0:nxp1, 1:nyp1, 0:nz] - Ez[0:nxp1, 0:ny, 0:nz]))

    # Hy update
    Hy[:,:,:] = (Chyh * Hy[:,:,:] + 
        Chyez * (Ez[1:nxp1, 0:nyp1, 0:nz] - Ez[0:nx, 0:nyp1, 0:nz]) + 
        Chyex * (Ex[0:nx, 0:nyp1, 1:nzp1] - Ex[0:nx, 0:nyp1, 0:nz]))
    # Hz update
    Hz[:,:,:] = (Chzh * Hz[:,:,:] + 
        Chzex * (Ex[0:nx, 1:nyp1, 0:nzp1] - Ex[0:nx, 0:ny, 0:nzp1]) + 
        Chzey * (Ey[1:nxp1, 0:ny, 0:nzp1] - Ey[0:nx, 0:ny, 0:nzp1]))
def update_magnetic_field_CPML_ABC():
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez
    from  calculate_domain_size import nxp1,ny,nz,nx,nyp1,nzp1
    # apply CPML to magnetic field components
    if is_cpml_xn:
        for i in np.arange(n_cpml_xn):
            Psi_hyx_xn[i,:,:] = cpml_b_mx_xn[i] * Psi_hyx_xn[i,:,:] \
                + cpml_a_mx_xn[i]*(Ez[i+1,:,:]-Ez[i,:,:]) 
            Psi_hzx_xn[i,:,:] = cpml_b_mx_xn[i] * Psi_hzx_xn[i,:,:] \
                + cpml_a_mx_xn[i]*(Ey[i+1,:,:]-Ey[i,:,:]) 
        Hy[0:n_cpml_xn,:,:] = Hy[0:n_cpml_xn,:,:] \
            + CPsi_hyx_xn[:,:,:] * Psi_hyx_xn[:,:,:]
        Hz[0:n_cpml_xn,:,:] = Hz[0:n_cpml_xn,:,:] \
            + CPsi_hzx_xn[:,:,:] * Psi_hzx_xn[:,:,:]

    if is_cpml_xp:
        n_st = nx - n_cpml_xp
        for i in np.arange(n_cpml_xp):
            Psi_hyx_xp[i,:,:] = cpml_b_mx_xp[i] * Psi_hyx_xp[i,:,:] \
                + cpml_a_mx_xp[i]*(Ez[i+n_st+1,:,:]-Ez[i+n_st,:,:]) 
            Psi_hzx_xp[i,:,:] = cpml_b_mx_xp[i] * Psi_hzx_xp[i,:,:] \
                + cpml_a_mx_xp[i]*(Ey[i+n_st+1,:,:]-Ey[i+n_st,:,:]) 

        Hy[n_st:nx,:,:] += CPsi_hyx_xp[:,:,:] * Psi_hyx_xp[:,:,:]
        Hz[n_st:nx,:,:] += CPsi_hzx_xp[:,:,:] * Psi_hzx_xp[:,:,:]    

    if is_cpml_yn:
        for i in np.arange(n_cpml_yn):
            Psi_hzy_yn[:,i,:] = cpml_b_my_yn[i] * Psi_hzy_yn[:,i,:] \
                + cpml_a_my_yn[i]*(Ex[:,i+1,:]-Ex[:,i,:]) 
            Psi_hxy_yn[:,i,:] = cpml_b_my_yn[i] * Psi_hxy_yn[:,i,:] \
                + cpml_a_my_yn[i]*(Ez[:,i+1,:]-Ez[:,i,:]) 
        Hz[:,0:n_cpml_yn,:] += CPsi_hzy_yn[:,:,:] * Psi_hzy_yn[:,:,:]
        Hx[:,0:n_cpml_yn,:] += CPsi_hxy_yn[:,:,:] * Psi_hxy_yn[:,:,:]    

    if is_cpml_yp:
        n_st = ny - n_cpml_yp
        for i in np.arange(n_cpml_yp):
            Psi_hzy_yp[:,i,:] = cpml_b_my_yp[i] * Psi_hzy_yp[:,i,:] \
                + cpml_a_my_yp[i]*(Ex[:,i+n_st+1,:]-Ex[:,i+n_st,:]) 
            Psi_hxy_yp[:,i,:] = cpml_b_my_yp[i] * Psi_hxy_yp[:,i,:] \
                + cpml_a_my_yp[i]*(Ez[:,i+n_st+1,:]-Ez[:,i+n_st,:]) 

        Hz[:,n_st:ny,:] += CPsi_hzy_yp[:,:,:] * Psi_hzy_yp[:,:,:]
        Hx[:,n_st:ny,:] += CPsi_hxy_yp[:,:,:] * Psi_hxy_yp[:,:,:]    

    if is_cpml_zn:
        for i in np.arange(n_cpml_zn):
            Psi_hxz_zn[:,:,i] = cpml_b_mz_zn[i] * Psi_hxz_zn[:,:,i] \
                + cpml_a_mz_zn[i]*(Ey[:,:,i+1]-Ey[:,:,i]) 
            Psi_hyz_zn[:,:,i] = cpml_b_mz_zn[i] * Psi_hyz_zn[:,:,i] \
                + cpml_a_mz_zn[i]*(Ex[:,:,i+1]-Ex[:,:,i]) 
        Hx[:,:,0:n_cpml_zn] += CPsi_hxz_zn[:,:,:] * Psi_hxz_zn[:,:,:]
        Hy[:,:,0:n_cpml_zn] += CPsi_hyz_zn[:,:,:] * Psi_hyz_zn[:,:,:]    

    if is_cpml_zp:
        n_st = nz - n_cpml_zp
        for i in np.arange(n_cpml_zp):
            Psi_hxz_zp[:,:,i] = cpml_b_mz_zp[i] * Psi_hxz_zp[:,:,i] \
                + cpml_a_mz_zp[i]*(Ey[:,:,i+n_st+1]-Ey[:,:,i+n_st]) 
            Psi_hyz_zp[:,:,i] = cpml_b_mz_zp[i] * Psi_hyz_zp[:,:,i] \
                + cpml_a_mz_zp[i]*(Ex[:,:,i+n_st+1]-Ex[:,:,i+n_st]) 
        Hx[:,:,n_st:nz] += CPsi_hxz_zp[:,:,:] * Psi_hxz_zp[:,:,:]
        Hy[:,:,n_st:nz] += CPsi_hyz_zp[:,:,:] * Psi_hyz_zp[:,:,:]    

def update_electric_fields():
    from run_fdtd_time_marching_loop import current_time
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez,E_x,E_y,E_z
    from initialize_dispersive_coefficients import Qxp,Cqqx,Cqex,Qyp,Cqqy,Cqey,Qzp,Cqqz,Cqez
    from initialize_updating_coefficients import Cexe,Ceye,Ceze
    from initialize_fdtd_material_grid import eps_r_x,material_3d_space,deps_x,tau_k_x,eps_r_y,deps_y,tau_k_y,tau_k_z
    from initialize_updating_coefficients import psi_x,psi_y,psi_z,den_x,den_y,den_z
    current_time  = current_time + dt/2
    # Regular field updates
    E_x[0:nx,1:ny,1:nz]=Ex[0:nx,1:ny,1:nz]
    Ex[0:nx,1:ny,1:nz] = Cexe[0:nx,1:ny,1:nz]*Ex[0:nx,1:ny,1:nz] \
                        + Cexhz[0:nx,1:ny,1:nz]*\
                        (Hz[0:nx,1:ny,1:nz]-Hz[0:nx,0:ny-1,1:nz]) \
                        + Cexhy[0:nx,1:ny,1:nz]*\
                        (Hy[0:nx,1:ny,1:nz]-Hy[0:nx,1:ny,0:nz-1])\
                        - (dy) * Cexhz[0:nx,1:ny,1:nz] * (Qxp[0:nx,1:ny,1:nz])*2*tau_k_x[0:nx,1:ny,1:nz]/den_x[0:nx,1:ny,1:nz]
    Qxp[0:nx,1:ny,1:nz] = \
                Cqqx[0:nx,1:ny,1:nz] * Qxp[0:nx,1:ny,1:nz]\
            + Cqex[0:nx,1:ny,1:nz]  * (Ex[0:nx,1:ny,1:nz]-E_x[0:nx,1:ny,1:nz])

    E_y[1:nx,0:ny,1:nz]=Ey[1:nx,0:ny,1:nz]
    Ey[1:nx,0:ny,1:nz]=Ceye[1:nx,0:ny,1:nz]*Ey[1:nx,0:ny,1:nz] \
                        + Ceyhx[1:nx,0:ny,1:nz]*  \
                        (Hx[1:nx,0:ny,1:nz]-Hx[1:nx,0:ny,0:nz-1]) \
                        + Ceyhz[1:nx,0:ny,1:nz]*  \
                        (Hz[1:nx,0:ny,1:nz]-Hz[0:nx-1,0:ny,1:nz])\
                        - (dz) * Ceyhx[1:nx,0:ny,1:nz] * (Qyp[1:nx,0:ny,1:nz])*2*tau_k_y[1:nx,0:ny,1:nz]/den_y[1:nx,0:ny,1:nz]
    Qyp[1:nx,0:ny,1:nz] = \
                Cqqy[1:nx,0:ny,1:nz]  * Qyp[1:nx,0:ny,1:nz] \
            + Cqey[1:nx,0:ny,1:nz]  * (Ey[1:nx,0:ny,1:nz]-E_y[1:nx,0:ny,1:nz])  
    E_z[1:nx,1:ny,0:nz]=Ez[1:nx,1:ny,0:nz]
    #if(time_step==1):print('Ez1:',Ez[20,20,20])
    Ez[1:nx,1:ny,0:nz]=Ceze[1:nx,1:ny,0:nz]*Ez[1:nx,1:ny,0:nz] \
                        + Cezhy[1:nx,1:ny,0:nz]*  \
                        (Hy[1:nx,1:ny,0:nz]-Hy[0:nx-1,1:ny,0:nz]) \
                        + Cezhx[1:nx,1:ny,0:nz]*\
                        (Hx[1:nx,1:ny,0:nz]-Hx[1:nx,0:ny-1,0:nz])\
                        - (dx) * Cezhy[1:nx,1:ny,0:nz] * (Qzp[1:nx,1:ny,0:nz])*2*tau_k_z[1:nx,1:ny,0:nz]/den_z[1:nx,1:ny,0:nz]
    #if(time_step==1):print('Ez2:',Ez[20,20,20])  #time_step=2  Ez2=-1.7029e-012
    Qzp[1:nx,1:ny,0:nz] = \
                Cqqz[1:nx,1:ny,0:nz]  * Qzp[1:nx,1:ny,0:nz] \
            + Cqez[1:nx,1:ny,0:nz] * (Ez[1:nx,1:ny,0:nz]-E_z[1:nx,1:ny,0:nz])

def update_electric_field_CPML_ABC():
    # apply CPML to electric field components
    from initialize_boundary_conditions import n_cpml_xn
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez,E_x,E_y,E_z
    if is_cpml_xn:
        for i in np.arange(n_cpml_xn):
            Psi_eyx_xn[i,:,:] = cpml_b_ex_xn[i] * Psi_eyx_xn[i,:,:] \
                + cpml_a_ex_xn[i]*(Hz[i+1,:,:]-Hz[i,:,:]) 
            Psi_ezx_xn[i,:,:] = cpml_b_ex_xn[i] * Psi_ezx_xn[i,:,:] \
                + cpml_a_ex_xn[i]*(Hy[i+1,:,:]-Hy[i,:,:]) 
        Ey[1:n_cpml_xn+1,:,:] += CPsi_eyx_xn * Psi_eyx_xn
        Ez[1:n_cpml_xn+1,:,:] += CPsi_ezx_xn * Psi_ezx_xn    

    if is_cpml_xp :
        n_st = nx - n_cpml_xp
        for i in np.arange (n_cpml_xp):
            Psi_eyx_xp[i,:,:] = cpml_b_ex_xp[i] * Psi_eyx_xp[i,:,:] \
                + cpml_a_ex_xp[i]*(Hz[i+n_st,:,:]-Hz[i+n_st-1,:,:]) 
            Psi_ezx_xp[i,:,:] = cpml_b_ex_xp[i] * Psi_ezx_xp[i,:,:] \
                + cpml_a_ex_xp[i]*(Hy[i+n_st,:,:]-Hy[i+n_st-1,:,:]) 
                
        Ey[n_st:nx,:,:] += CPsi_eyx_xp * Psi_eyx_xp
        Ez[n_st:nx,:,:] += CPsi_ezx_xp * Psi_ezx_xp

    if is_cpml_yn:
        for i in np.arange(n_cpml_yn):
            Psi_ezy_yn[:,i,:] = cpml_b_ey_yn[i] * Psi_ezy_yn[:,i,:] \
                + cpml_a_ey_yn[i]*(Hx[:,i+1,:]-Hx[:,i,:]) 
            Psi_exy_yn[:,i,:] = cpml_b_ey_yn[i] * Psi_exy_yn[:,i,:] \
                + cpml_a_ey_yn[i]*(Hz[:,i+1,:]-Hz[:,i,:]) 
        Ez[:,1:n_cpml_yn+1,:] += CPsi_ezy_yn * Psi_ezy_yn
        Ex[:,1:n_cpml_yn+1,:] += CPsi_exy_yn * Psi_exy_yn    

    if is_cpml_yp:
        n_st = ny - n_cpml_yp
        for i in np.arange(n_cpml_yp):
            Psi_ezy_yp[:,i,:] = cpml_b_ey_yp[i] * Psi_ezy_yp[:,i,:] \
                + cpml_a_ey_yp[i]*(Hx[:,i+n_st,:]-Hx[:,i+n_st-1,:]) 
            Psi_exy_yp[:,i,:] = cpml_b_ey_yp[i] * Psi_exy_yp[:,i,:] \
                + cpml_a_ey_yp[i]*(Hz[:,i+n_st,:]-Hz[:,i+n_st-1,:]) 
                
        Ez[:,n_st:ny,:] += CPsi_ezy_yp * Psi_ezy_yp
        Ex[:,n_st:ny,:] += CPsi_exy_yp * Psi_exy_yp    

    if is_cpml_zn:
        for i in np.arange(n_cpml_zn):
            Psi_exz_zn[:,:,i] = cpml_b_ez_zn[i] * Psi_exz_zn[:,:,i] \
                + cpml_a_ez_zn[i]*(Hy[:,:,i+1]-Hy[:,:,i]) 
            Psi_eyz_zn[:,:,i] = cpml_b_ez_zn[i] * Psi_eyz_zn[:,:,i] \
                + cpml_a_ez_zn[i]*(Hx[:,:,i+1]-Hx[:,:,i]) 
                
        Ex[:,:,1:n_cpml_zn+1] += CPsi_exz_zn * Psi_exz_zn
        Ey[:,:,1:n_cpml_zn+1] += CPsi_eyz_zn * Psi_eyz_zn    

    if is_cpml_zp:
        n_st = nz - n_cpml_zp
        for i in np.arange(n_cpml_zp):
            Psi_exz_zp[:,:,i] = cpml_b_ez_zp[i] * Psi_exz_zp[:,:,i] \
                + cpml_a_ez_zp[i]*(Hy[:,:,i+n_st]-Hy[:,:,i+n_st-1]) 
            Psi_eyz_zp[:,:,i] = cpml_b_ez_zp[i] * Psi_eyz_zp[:,:,i] \
                + cpml_a_ez_zp[i]*(Hx[:,:,i+n_st]-Hx[:,:,i+n_st-1]) 
                
        Ex[:,:,n_st:nz] += CPsi_exz_zp * Psi_exz_zp
        Ey[:,:,n_st:nz] += CPsi_eyz_zp * Psi_eyz_zp

def update_voltage_sources():
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez
    from define_output_parameters import source
    from define_problem_space_parameters import waveform_type,f_sinwave,f_center,setbandwidth
    # print(Ez[source["x"]-1,source["y"]-1,source["z"]-1])
    # print(source["x"],source["y"],source["z"]) #83 44 44
    # exit(-1)
    #4:0.756257
    if waveform_type ==1:
        Ez[source["x"],source["y"],source["z"]]+= 10e-5*dt/(eps_0*dx*dy*dz)*np.sin(2*np.pi*f_sinwave*dt*(time_step+1))
    else:
        tau = 0.966/setbandwidth
        t_0 = 4.5 * tau
        Ez[source["x"],source["y"],source["z"]]+=np.cos(2*np.pi*f_center*(dt*(time_step+1)-t_0/2))*np.exp(-(dt*(time_step+1)-t_0/2)**2/tau**2)
def update_current_sources():
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez
    from define_sources_and_lumped_elements import current_sources,number_of_current_sources
    # updating electric field components 
    # associated with the current sources
    for ind in np.arange(number_of_current_sources):
        fi = current_sources[ind]["field_indices"]
        cond_string = current_sources[ind]["direction"][0]
        if cond_string == 'x' :
            Ex[fi] += current_sources[ind]["Cexs"] \
                * current_sources[ind]["current_per_e_field"][time_step]
        elif cond_string == 'y':
            Ey[fi] += current_sources[ind]["Ceys"] \
                * current_sources[ind]["current_per_e_field"][time_step]
        elif cond_string == 'z':
            Ez[fi] += current_sources[ind]["Cezs"] \
                * current_sources[ind]["current_per_e_field"][time_step]
def capture_sampled_electric_fields():
    # Capturing electric fields
    from initialize_output_parameters import number_of_sampled_electric_fields,sampled_electric_fields
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez
    for ind in np.arange(number_of_sampled_electric_fields):
        is0 = sampled_electric_fields[ind]["is"]
        js = sampled_electric_fields[ind]["js"]
        ks = sampled_electric_fields[ind]["ks"]
        str_component = sampled_electric_fields[ind]["component"]
        if str_component == 'x':
            sampled_value = 0.5 * np.sum(Ex[is0-2:is0,js-1,ks-1]) #Ex(is0,js,ks)#
        elif str_component == 'y':
            sampled_value = 0.5 * np.sum(Ey[is0-1,js-2:js,ks-1]) #Ey(is0,js,ks)#
        elif str_component == 'z':
            sampled_value = 0.5*np.sum(Ez[is0-1,js-1,ks-2:ks]) #Ez(is0,js,ks)#0.5 * 
        elif str_component == 'm':
            svx = 0.5 * np.sum(Ex[is0-2:is0,js-1,ks-1]) 
            svy = 0.5 * np.sum(Ey[is0-1,js-2:js,ks-1]) 
            svz = 0.5 * np.sum(Ez[is0-1,js-1,ks-2:ks]) 
            sampled_value = np.sqrt(svx**2 + svy**2 + svz**2)
        sampled_electric_fields[ind]["sampled_value"][time_step] = sampled_value

def capture_sampled_magnetic_fields():
    # Capturing magnetic fields
    from initialize_output_parameters import number_of_sampled_magnetic_fields,sampled_magnetic_fields
    from initialize_fdtd_parameters_and_arrays import Hx,Hy,Hz,Ex,Ey,Ez
    for ind in np.arange(number_of_sampled_magnetic_fields):
        is0 = sampled_magnetic_fields[ind]["is"]
        js = sampled_magnetic_fields[ind]["js"]
        ks = sampled_magnetic_fields[ind]["ks"]
        str_component = sampled_magnetic_fields[ind]["component"]
        if str_component == 'x':
            sampled_value = 0.25 * np.sum(Hx[is0-1,js-2:js,ks-2:ks]) 
        elif str_component == 'y':
            sampled_value = 0.25 * np.sum(Hy[is0-2:is0,js-1,ks-2:ks]) 
        elif str_component == 'z':
            sampled_value = 0.25 * np.sum(Hz[is0-2:is0-1,js-2:js-1,ks-1]) 
        elif str_component == 'm':
            svx = 0.25 * np.sum(Hx[is0-1,js-2:js,ks-2:ks]) 
            svy = 0.25 * np.sum(Hy[is0-2:is0,js-1,ks-2:ks]) 
            svz = 0.25 * np.sum(Hz[is0-2:is0-1,js-2:js-1,ks-1]) 
            sampled_value = np.sqrt(svx**2 + svy**2 + svz**2)
        sampled_magnetic_fields[ind]["sampled_value"][time_step] = sampled_value
from fileutils import ElaspedTime

et = ElaspedTime()
et.elapsed()
for time_step in np.arange(number_of_time_steps):
    print(f'time_step={time_step+1}')
    if(time_step%20==1):
        et.elapsed()
    from debug import set_time_step
    set_time_step(time_step)
    update_magnetic_fields()
    update_magnetic_field_CPML_ABC()
#     
#     capture_sampled_currents
    update_electric_fields()
    update_electric_field_CPML_ABC()
    update_voltage_sources()
    #update_current_sources()
#     update_inductors 
#     update_diodes
    capture_sampled_electric_fields()
    capture_sampled_magnetic_fields()
#     capture_sampled_voltages
#     printlay_sampled_parameters
et.elapsed()
end_time = time.process_time()
total_time_in_minutes = (end_time - start_time)/60
print(f'Total simulation time is {total_time_in_minutes} minutes.')
